from torch.utils.data import DataLoader,Dataset
import torch
import numpy as np
import cv2
import os
import pandas as pd
from data_aug import resize,data_aug,random_rotate_and_crop
from process_data import get_guanz,conver_csv,convert_dict
import torch.nn.functional as F


def make_single_heatmap(img_width, img_height, c_x, c_y, sigma):
    ## 如果坐标为0，则表示该点实际为缺失点。生成空白的heatmap
    if c_x == 0 or c_y == 0:
        sigma = 1
    X1 = np.linspace(1, img_width, img_width)
    Y1 = np.linspace(1, img_height, img_height)
    [X, Y] = np.meshgrid(X1, Y1)
    X = X - c_x
    Y = Y - c_y
    D2 = X * X + Y * Y
    E2 = 2.0 * sigma * sigma
    Exponent = D2 / E2
    heatmap = np.exp(-Exponent)
    return 255 * heatmap[np.newaxis, ...]


def make_multi_heatmap(size,anno_list,sigma = 5):
    ## 生成每个关键点对应一层heatmap的形式。
    ## 传入按照固定顺序排列的坐标的list
    heatmap_list = []
    for i in anno_list:
        single_heatmap = make_single_heatmap(img_width=size[0],img_height=size[1],
                                                  c_x=i[0],c_y=i[1],sigma=sigma)
        heatmap_list.append(single_heatmap)
    heatmap = np.concatenate(heatmap_list,axis=0)
    return heatmap


def make_heatmap(size,anno_list,sigma = 5):
    ## 传入按照固定顺序排列的坐标的list

    ## 将点分为锥体和椎间两类，分来生成heatmap
    anno_list_1 = anno_list[:5]
    anno_list_2 = anno_list[5:]

    heatmap_1 = make_single_heatmap(img_width=size[0],img_height=size[1],
                                    c_x=anno_list_1[0][0],c_y=anno_list_1[0][1],sigma=sigma)
    for i in anno_list_1[1:]:
        heatmap_1 = heatmap_1+make_single_heatmap(img_width=size[0],img_height=size[1],
                                                  c_x=i[0],c_y=i[1],sigma=sigma)


    heatmap_2 = make_single_heatmap(img_width=size[0], img_height=size[1],
                                    c_x=anno_list_2[0][0], c_y=anno_list_2[0][1], sigma=sigma)
    for i in anno_list_2[1:]:
        heatmap_2 = heatmap_2 + make_single_heatmap(img_width=size[0], img_height=size[1],
                                                    c_x=i[0], c_y=i[1], sigma=sigma)
    heatmap = np.concatenate((heatmap_1,heatmap_2),0)

    # # print(heatmap_1.shape,heatmap_2.shape)
    # cv2.imshow('heatmap_1',heatmap_1[0])
    # cv2.waitKey()
    # cv2.imshow('heatmap_2',heatmap_2[0])
    # cv2.waitKey()

    return heatmap

def heatmap_to_anno(heatmap):
    assert heatmap.shape[0] == 2,'the channel of heatmap should be 2'
    def heatmap_nms(heat,kernel = 13):
        pad = (kernel - 1 )//2
        hmax = F.max_pool2d(
            heat[np.newaxis,...],(kernel,kernel),stride=1,padding = pad
        )
        keep = (hmax == heat).float()
        return (heat * keep)[0]

    def get_top_K(heatmap,k=1):
        anno_list = []
        heatmap = heatmap_nms(heatmap)
        heatmap_view = heatmap.view(-1)
        top_k_num = torch.topk(heatmap_view, k)[0]
        heatmap = heatmap.cpu()
        anno = np.where(heatmap >= top_k_num[-1].cpu())
        for i,j in zip(anno[0][:k],anno[1][:k]):
            anno_list.append((i+1,j+1))
        ## 按照位置高低进行排序
        anno_list = sorted(anno_list,key=lambda j:j[0])
        ## 将坐标转成坐标系中的形式(x,y)
        anno_list = [(j[1],j[0])for j in anno_list]
        return anno_list

    anno_list_1 = get_top_K(heatmap[0],k=5)
    anno_list_2 = get_top_K(heatmap[1],k=6)
    return anno_list_1+anno_list_2

def multi_heatmap_to_anno(heatmap):
    assert heatmap.shape[0] == 11,'the channel of heatmap shoudl be 11'
    anno_list = []
    heatmap = heatmap.cpu().numpy()
    for index in range(heatmap.shape[0]):
        single_heatmap = heatmap[index]
        anno = np.where(single_heatmap == single_heatmap.max())
        ## 转成xy的形式
        anno_list.append((int(anno[1][0]),int(anno[0][0])))
    return anno_list


def make_heatmap_with_category(size,anno_list,category,sigma = 5):
    ## 传入按照固定顺序排列的坐标的list
    ## 将点分为锥体和椎间两类，分来生成heatmap
    anno_list_1 = anno_list[:5]
    anno_list_2 = anno_list[5:]
    ## 将类别转换成数字。从0开始，
    category_1 = [int('v2' in p) for p in category[:5]]
    category_2 = [int(p[1]) - 1 for p in category[5:]]


    heatmap_1 = np.zeros(shape=((len(category_1),size[0],size[1])))
    for i in range(len(anno_list_1)):
        heatmap_1[category_1[i]] = heatmap_1[category_1[i]] + make_single_heatmap(img_width=size[0], img_height=size[1],
                                                    c_x=anno_list_1[i][0], c_y=anno_list_1[i][1], sigma=sigma)

    heatmap_2 = np.zeros(shape=((len(category_2), size[0], size[1])))
    for i in range(len(anno_list_2)):
        heatmap_2[category_2[i]] = heatmap_2[category_2[i]] + make_single_heatmap(img_width=size[0], img_height=size[1],
                                                                                  c_x=anno_list_2[i][0],
                                                                                  c_y=anno_list_2[i][1], sigma=sigma)
    heatmap = np.concatenate((heatmap_1, heatmap_2), 0)

    # # # print(heatmap_1.shape,heatmap_2.shape)
    # cv2.imshow('heatmap_1', heatmap_1[1])
    # cv2.waitKey()
    # cv2.imshow('heatmap_2',heatmap_2[1])
    # cv2.waitKey()
    # print('**'*30)
    # print(heatmap.shape)

    return heatmap


def heatmap_to_anno_with_category(heatmap):
    heatmap_1 = torch.sum(heatmap[:5],dim=1)
    heatmap_2 = torch.sum(heatmap[5:],dim=1)

    def heatmap_nms(heat, kernel=13):
        pad = (kernel - 1) // 2
        hmax = F.max_pool2d(
            heat[np.newaxis, ...], (kernel, kernel), stride=1, padding=pad
        )
        keep = (hmax == heat).float()
        return (heat * keep)[0]

    def get_top_K(heatmap, k=1):
        anno_list = []
        heatmap = heatmap_nms(heatmap)
        heatmap_view = heatmap.view(-1)
        top_k_num = torch.topk(heatmap_view, k)[0]
        heatmap = heatmap.cpu()
        anno = np.where(heatmap >= top_k_num[-1].cpu())
        for i, j in zip(anno[0][:k], anno[1][:k]):
            anno_list.append((i + 1, j + 1))
        ## 按照位置高低进行排序
        anno_list = sorted(anno_list, key=lambda j: j[0])
        ## 将坐标转成坐标系中的形式(x,y)
        anno_list = [(j[1], j[0]) for j in anno_list]
        return anno_list

    anno_list_1 = get_top_K(heatmap_1, k=5)
    anno_list_2 = get_top_K(heatmap_2, k=6)
    return anno_list_1 + anno_list_2



class Mydataset(Dataset):
    def __init__(self,anno_list,size = (512,512),image_path=None,sigma = 3):
        self.anno_list = anno_list
        self.image_path = image_path
        self.size = size
        self.sigma = sigma
        self.processed_image,self.processed_label = self.process_data()


    def __getitem__(self, index):
        image = self.processed_image[index]
        anno = self.processed_label[index]
        image, anno = data_aug(image, anno)
        image, anno = resize(image, anno, self.size)
        # print(image_path,label)
        heatmap = make_heatmap(self.size, anno, self.sigma)
        # heatmap = make_multi_heatmap(self.size,anno,self.sigma)
        image, heatmap = random_rotate_and_crop(image, heatmap, p=0.6, crop_limit=20,angle_limit=15)
        image = torch.from_numpy(np.transpose(image, (2, 0, 1))).float()
        heatmap = torch.from_numpy(heatmap).float()
        return image,heatmap

    def __len__(self):
        return len(self.anno_list)

    def process_data(self):
        processed_image = []
        processed_label = []
        for index in range(len(self.anno_list)):
            try:
                image_path = self.anno_list[index]['image_path']
                image_path = os.path.join(self.image_path, image_path.split('//')[1] + '_' +
                                          image_path.split('//')[2].replace('dcm', 'jpg'))
                image = cv2.imread(image_path)
                # 得到某一图片对应的所有标签信息，包括坐标位置和每个点属性
                label = self.anno_list[index]
                ## 得到坐标位置按照固定顺序排列的list，其中包含坐标信息和类别信息。用来生成heatmap使用。
                label = convert_dict(label)
                anno = [p[0] for p in label]
                category = [p[1] for p in label]
                # image, anno = resize(image, anno, self.size)
                processed_image.append(image)
                processed_label.append(anno)
            except Exception:
                continue
        return processed_image,processed_label


if __name__ == '__main__':
    import time
    start_time = time.time()
    my_dataset = Mydataset(csv_file='../data/csv_label/train150_anno.csv',image_path='../data/image')
    dataloader = DataLoader(dataset=my_dataset,batch_size=4,num_workers=0)
    for batch_image,batch_heatmap in dataloader:
        pass
        print(batch_image.shape,batch_heatmap.shape)
        # print(cate1,cate2)
    print(time.time() - start_time)